"""
rerank_client.py
Silicon Flow 通用 rerank 封装（requests 实现）。
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass
from typing import List, Sequence

import requests
from dotenv import load_dotenv

load_dotenv()
logger = logging.getLogger(__name__)


# ------------------ 自定义异常 ------------------
class RerankError(RuntimeError):
    """任何 rerank 调用异常基类。"""


class HTTPError(RerankError):
    """HTTP 非 2xx。"""


# ------------------ 数据模型 ------------------
@dataclass(slots=True)
class RerankResult:
    """单条 rerank 结果，字段清晰可读。"""

    text: str
    index: int
    score: float

    def __str__(self) -> str:
        return f"RerankResult(text={self.text!r}, index={self.index}, score={self.score:.4f})"


# ------------------ 核心封装 ------------------
class RerankClient:
    """
    轻量级 Silicon Flow rerank 客户端（bge-reranker-v2-m3）。

    用法
    ----
    >>> client = RerankClient()
    >>> ranked = client.rerank("Apple", ["apple", "banana", "fruit"])
    >>> print(ranked[0])
    RerankResult(text='apple', index=0, score=0.9918)
    """

    def __init__(
        self,
        api_key: str | None = None,
        *,
        base_url: str = "https://api.siliconflow.cn/v1/rerank",
        model: str = "BAAI/bge-reranker-v2-m3",
        timeout: float = 30.0,
    ):
        self._api_key = api_key or os.getenv("SILICONCLOUD_API_KEY")
        if not self._api_key:
            raise ValueError("api_key 未提供且未找到环境变量 SILICONCLOUD_API_KEY")

        self._base_url = base_url.strip()
        self._model = model
        self._timeout = timeout
        self._session = requests.Session()
        self._session.headers.update(
            {
                "Authorization": f"Bearer {self._api_key}",
                "Content-Type": "application/json",
            }
        )

    def rerank(
        self,
        query: str,
        documents: Sequence[str],
        *,
        top_k: int | None = None,
        return_documents: bool = False,
        max_chunks_per_doc: int | None = None,
        overlap_tokens: int | None = None,
    ) -> List[RerankResult]:
        """
        对给定 query 与文档列表进行相关性重排。

        Parameters
        ----------
        query : str
            查询文本。
        documents : Sequence[str]
            候选文档，长度 ≤ 2048（模型限制）。
        top_k : int, optional
            仅返回前 k 条结果；默认全部返回。
        return_documents : bool, default=False
            如果为False，响应不包含文档文本；如果为True，包含输入文档文本。
        max_chunks_per_doc : int, optional
            每个文档生成的最大块数。长文档会被分成多个块进行计算，
            取块中最高分作为文档的分数。仅 BAAI/bge-reranker-v2-m3, 
            Pro/BAAI/bge-reranker-v2-m3, netease-youdao/bce-reranker-base_v1 支持此字段。
        overlap_tokens : int, optional
            文档分块时相邻块之间的令牌重叠数量。仅 BAAI/bge-reranker-v2-m3, 
            Pro/BAAI/bge-reranker-v2-m3, netease-youdao/bce-reranker-base_v1 支持此字段。
            取值范围: ≤ 80。

        Returns
        -------
        List[RerankResult]
            已按 score 降序排序，score 保留 4 位小数。
        """
        payload = {
            "model": self._model,
            "query": query,
            "documents": list(documents),
            "top_n": top_k,  # API使用top_n参数，我们保持top_k接口兼容性
            "return_documents": return_documents,
        }
        
        # 仅当提供了值时才添加可选参数
        if max_chunks_per_doc is not None:
            payload["max_chunks_per_doc"] = max_chunks_per_doc
            
        if overlap_tokens is not None:
            if overlap_tokens > 80:
                logger.warning("overlap_tokens超过80的限制，已自动调整为80")
                overlap_tokens = 80
            payload["overlap_tokens"] = overlap_tokens

        try:
            resp = self._session.post(
                self._base_url, json=payload, timeout=self._timeout
            )
            resp.raise_for_status()
        except requests.HTTPError as e:
            raise HTTPError(
                f"HTTP {e.response.status_code}: {e.response.text}"
            ) from e
        except requests.RequestException as e:
            raise RerankError("Network error") from e

        data = resp.json()
        results = data.get("results", [])
        if not results:
            raise RerankError("Unexpected response format")

        # 构建 dataclass 列表，并一次性按 score 降序
        ranked = [
            RerankResult(
                text=documents[item["index"]],
                index=item["index"],
                score=round(item["relevance_score"], 4),
            )
            for item in results
        ]
        ranked.sort(key=lambda r: r.score, reverse=True)

        logger.debug("rerank success, %d results returned", len(ranked))
        return ranked